from patagent.task import execute_query, dim_table_report_chain
import json


def query(state):
    """
    Query database

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---QUERY DATABASE---")
    sql = state["sql"]

    data = execute_query(sql)
    return {"sql": sql, "data": json.dumps(data)}


def generate_report(state):
    """
    generate_report answer

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, generation, that contains LLM generation
    """
    print("---generate_report---")
    sql = state["sql"]
    data = state["data"]

    # RAG generation
    generation = dim_table_report_chain.invoke({"question": data})
    return {"sql": sql, "data": data, "generation": generation}



from langgraph.graph import END, StateGraph, START
from patagent.agent.patent_rag import *
from typing_extensions import TypedDict


class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        sql: query sql
        data: data from database
        generation: LLM generation
    """
    sql: str
    data: str
    generation: str

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("query", query)  # query
workflow.add_node("generate_report", generate_report)  # generatae

# Build graph
workflow.add_edge(START, "query")
workflow.add_edge("query", "generate_report")
workflow.add_edge("generate_report", END)

# Compile
graph = workflow.compile()

# print(graph.invoke({'sql': "select dt,table_name,dim,insert_cnt,update_cnt from quality_dim_count where table_name in ('dwd_llm_paper') and dt >= DATE_SUB(CURDATE(), INTERVAL 15 DAY) ORDER BY dt DESC"})['generation'])